import os
import pandas as pd

def attach_country_attributes(result_df, output_path, data_dir):
    """
    给结果文件附加国家属性数据，并保存新的 CSV 文件

    参数：
        result_df (pd.DataFrame): 排序结果 DataFrame（需包含 '国名Ch' 列）
        output_path (str): 原始结果 CSV 文件路径
        data_dir (str): 国家属性文件所在目录，应包含 'countryAttr.csv'

    返回：
        enriched_df (pd.DataFrame or None): 成功附加后返回新的 DataFrame，失败返回 None
    """
    attr_file = os.path.join(data_dir, "countryAttr.csv")
    required_cols = ['国名Ch', '地理区域', '收入水平', '宗教信仰']

    if not os.path.exists(attr_file):
        print("⚠ 未找到国家属性文件 countryAttr.csv，已跳过附加操作")
        return None

    enrich = input("\n🧩 是否附加国家属性（地理区域、收入水平、宗教信仰等）？Y/N：").strip().lower()
    if enrich != 'y':
        print("✅ 已跳过附加国家属性步骤")
        return None

    try:
        attr_df = pd.read_csv(attr_file)

        # ✅ 检查字段完整性
        missing_cols = [col for col in required_cols if col not in attr_df.columns]
        if missing_cols:
            print(f"❌ 错误：以下必要列在 {attr_file} 中缺失：{missing_cols}")
            return None

        before_rows = result_df.shape[0]
        enriched_df = result_df.merge(attr_df, on='国名Ch', how='left', suffixes=('', '_attr'))

        # ✅ 移除重复的“序号”列（若存在）
        if '序号_attr' in enriched_df.columns:
            enriched_df.drop(columns=['序号_attr'], inplace=True)

        after_rows = enriched_df.shape[0]
        if after_rows < before_rows:
            print("⚠️ 警告：合并后数据行数减少，可能存在“国名Ch”匹配失败的情况。请检查国家名称是否一致。")

        enriched_output = output_path.replace(".csv", "_附加属性.csv")
        enriched_df.to_csv(enriched_output, index=False, encoding='utf-8-sig')
        print(f"📌 含属性的排序结果已保存为：{enriched_output}")

        return enriched_df

    except Exception as e:
        print(f"❌ 附加国家属性过程中发生错误：{e}")
        return None
